#include "tunnel.h"

#include "muduo/base/Logging.h"
#include "muduo/net/Callbacks.h"
/**
 * @brief Construct a new Tunnel:: Tunnel object
 *
 * @param loop 事件循环
 * @param serverAddr 服务端的地址
 * @param servercoon 和客户端的一次TCP连接
 */
Tunnel::Tunnel(EventLoop *loop, const InetAddress &serverAddr, const TcpConnectionPtr &servercoon)
    : client_(loop, serverAddr, servercoon->name()), serverConn_(servercoon)
{
  LOG_INFO << "Tunnel " << servercoon->peerAddress().toIpPort() << " <-> " << serverAddr.toIpPort();
}
void Tunnel::setup()
{
  using std::placeholders::_1;
  using std::placeholders::_2;
  using std::placeholders::_3;
  client_.setConnectionCallback(std::bind(&Tunnel::onClientConnection, shared_from_this(), _1));

  client_.setMessageCallback(std::bind(&Tunnel::onClientMessage, shared_from_this(), _1, _2, _3));
  // 接收端设置高水位回调，限制1MB
  serverConn_->setHighWaterMarkCallback(
      std::bind(&Tunnel::onHighWaterMarkWeak, std::weak_ptr<Tunnel>(shared_from_this()),
                ServerClient::kServer, _1, _2),
      1 << 20);
}

void Tunnel::connect() { client_.connect(); }

void Tunnel::disconnect()
{
  client_.disconnect();
  // serverConn_.reset();
  teardown();
}

/**
 * @brief  客户端关闭连接, 重置和server的连接
 *
 */
void Tunnel::teardown()
{
  client_.setConnectionCallback(muduo::net::defaultConnectionCallback);
  client_.setMessageCallback(muduo::net::defaultMessageCallback);
  if (serverConn_) {
    serverConn_->setContext(boost::any());
    serverConn_->shutdown();
  }
  clientConn_.reset();
}
/**
 * @brief 在客户端上的连接状态更新了（连接上server了），设置一些基础信息
   1. 设置发送端（clien_）的高水位回调
   2. 给severConn_, 即原始客户端设置setContext，方便之后proxy可以直接send
 *
 * @param conn
 */
void Tunnel::onClientConnection(const muduo::net::TcpConnectionPtr &conn)
{
  LOG_DEBUG << (conn->connected() ? "UP" : "DOWN");
  using std::placeholders::_1;
  using std::placeholders::_2;
  if (conn->connected()) {
    conn->setTcpNoDelay(true);
    // 对发送端设置高水位回调，限制1MB
    conn->setHighWaterMarkCallback(
        std::bind(&Tunnel::onHighWaterMarkWeak, std::weak_ptr<Tunnel>(shared_from_this()),
                  ServerClient::kClient, _1, _2),
        1 << 20);
    // 设置关联，这样我们就可以通过getContext得到对应的client
    serverConn_->setContext(conn);
    serverConn_->startRead();
    clientConn_ = conn;
    // 发送已经积压的数据
    if (serverConn_->inputBuffer()->readableBytes() > 0) {
      conn->send(serverConn_->inputBuffer());
    }
  } else {
    teardown();
  }
}
/**
 * @brief 处理收到了客户端的信息
 *
 * @param conn 与客户端的一次tcp连接
 * @param buf  客户端发来的数据
 */
void Tunnel::onClientMessage(const muduo::net::TcpConnectionPtr &conn, muduo::net::Buffer *buf,
                             muduo::Timestamp)
{
  LOG_DEBUG << conn->name() << " " << buf->readableBytes();
  if (serverConn_ != nullptr)
    serverConn_->send(buf);
  else {  //这种情况不应该发生
    buf->retrieveAll();
    abort();
  }
}

/**
 * @brief 判断一个这个通道是否还在，是则启动高水位处理函数
 *
 * @param wkTunnel  TCP通道
 * @param which     发起的连接端
 * @param conn      这次TCP连接
 * @param bytesToSent  需要发送的字节数
 */
void Tunnel::onHighWaterMarkWeak(const std::weak_ptr<Tunnel> &wkTunnel, ServerClient which,
                                 const muduo::net::TcpConnectionPtr &conn, size_t bytesToSent)
{
  auto tunnel = wkTunnel.lock();
  if (tunnel != nullptr) tunnel->onHighWaterMark(which, conn, bytesToSent);
}

/**
 * @brief
 * 真正的高水位回调函数，具体思路为，高水位暂停读且另一端设定setWriteCompleteCallback
 *
 * @param which 用于判断是 client 还是 server 调用
 * @param conn 用于 LOG 写入 conn 的 name
 * @param bytesToSent 需要发送的字节数
 */
void Tunnel::onHighWaterMark(ServerClient which, const muduo::net::TcpConnectionPtr &conn,
                             size_t bytesToSent)
{
  using std::placeholders::_1;
  LOG_INFO << (which == ServerClient::kServer ? "server" : "client") << " onHighWaterMark "
           << conn->name() << " bytes " << bytesToSent;
  // 服务端输出的超过高水位了
  if (which == ServerClient::kServer) {
    if (serverConn_->outputBuffer()->readableBytes() > 0) {
      clientConn_->stopRead();
      serverConn_->setWriteCompleteCallback(std::bind(&Tunnel::onWriteCompleteWeak,
                                                      std::weak_ptr<Tunnel>(shared_from_this()),
                                                      ServerClient::kServer, _1));
    }
  } else {
    if (clientConn_->outputBuffer()->readableBytes() > 0) {
      serverConn_->stopRead();
      clientConn_->setWriteCompleteCallback(std::bind(&Tunnel::onWriteCompleteWeak,
                                                      std::weak_ptr<Tunnel>(shared_from_this()),
                                                      ServerClient::kClient, _1));
    }
  }
}

/**
 * @brief 判断一个这个通道是否还在，是则启动低（空）水位处理函数
 *
 * @param wkTunnel TCP通道
 * @param which 发起的连接端
 * @param conn 这次TCP连接
 */
void Tunnel::onWriteCompleteWeak(const std::weak_ptr<Tunnel> &wkTunnel, ServerClient which,
                                 const muduo::net::TcpConnectionPtr &conn)
{
  auto tunnel = wkTunnel.lock();
  if (tunnel != nullptr) tunnel->onWriteComplete(which, conn);
}

/**
 * @brief 真正的低（空）水位回调函数
 *
 * @param which 发起的连接端
 * @param conn 这次TCP连接
 */
void Tunnel::onWriteComplete(ServerClient which, const muduo::net::TcpConnectionPtr &conn)
{
  LOG_INFO << (which == ServerClient::kServer ? "server" : "client") << " onWriteComplete "
           << conn->name();
  // 如果 server 处理完了它的数据
  if (which == ServerClient::kServer) {
    clientConn_->startRead();
    // 清空低水位回调函数，目的是防止一直在没有数据的时候一直在回调
    serverConn_->setWriteCompleteCallback(muduo::net::WriteCompleteCallback{});
  } else {
    serverConn_->startRead();
    clientConn_->setWriteCompleteCallback(muduo::net::WriteCompleteCallback{});
  }
}

Tunnel::~Tunnel() { LOG_INFO << "~Tunnel"; }
